外部ネットワークアクセスで AWS サービスに対する IAM 認証を試してみた #SnowflakeDB

import json

def lambda_handler(event, context):
    # eventからnameの値を取得
    name = event.get('name')

    # nameがある場合は "Hello, name" の形式で、ない場合は "Hello!" を返す
    greeting = f"Hello, {name}" if name else "Hello!"

    return {
        'statusCode': 200,
        'body': json.dumps(greeting)
    }

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": "sts:AssumeRole",
            "Principal": {
                "AWS": "<自身のAWSアカウントID>"
            },
            "Condition": {
                "StringEquals": {
                    "sts:ExternalId": "<任意の値>"
                }
            }
        }
    ]
}

{
    "Version": "2012-10-17",
    "Statement":
    [
        {
        "Effect": "Allow",
        "Principal":
            {
            "AWS": "arn:aws:sts::xxxxxxxxxxxx:assumed-role/<IAMロール名>/snowflake"
            },
        "Action": "execute-api:Invoke",
        "Resource": "arn:aws:execute-api:ap-northeast-1:xxxxxxxxxxxx:xxxxxx/*/POST/hello"
        }
    ]
}

CREATE OR REPLACE NETWORK RULE aws_api_gateway_network_rule
  MODE = EGRESS
  TYPE = HOST_PORT
  VALUE_LIST = ('{api_id}.execute-api.ap-northeast-1.amazonaws.com');

CREATE OR REPLACE SECURITY INTEGRATION aws_api_gateway_security_integration
  TYPE = API_AUTHENTICATION
  AUTH_TYPE = AWS_IAM
  ENABLED = TRUE
  AWS_ROLE_ARN = 'arn:aws:iam::xxxxxxxxxxxx:role/snowflake-api-gateway-role';

DESC SECURITY INTEGRATION aws_api_gateway_security_integration;

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": "sts:AssumeRole",
            "Principal": {
                "AWS": "<API_AWS_IAM_USER_ARN>"
            },
            "Condition": {
                "StringEquals": {
                    "sts:ExternalId": "<API_AWS_EXTERNAL_ID>"
                }
            }
        }
    ]
}

CREATE OR REPLACE SECRET aws_api_gateway_access_token
  TYPE = CLOUD_PROVIDER_TOKEN
  API_AUTHENTICATION = aws_api_gateway_security_integration;

CREATE OR REPLACE EXTERNAL ACCESS INTEGRATION aws_api_gateway_external_access_integration
  ALLOWED_NETWORK_RULES = (aws_api_gateway_network_rule)
  ALLOWED_AUTHENTICATION_SECRETS = (aws_api_gateway_access_token)
  ENABLED = TRUE
  COMMENT = 'Test API Gateway connectivity';

CREATE OR REPLACE FUNCTION aws_api_gateway_post_function(name STRING)
  RETURNS VARCHAR
  LANGUAGE PYTHON
  EXTERNAL_ACCESS_INTEGRATIONS = (aws_api_gateway_external_access_integration)
  RUNTIME_VERSION = '3.11'
  SECRETS = ('cred' = <db>.<schema>.aws_api_gateway_access_token)
  PACKAGES = ('requests', 'aws-requests-auth')
  HANDLER = 'main_handler'
AS
$$
from aws_requests_auth.aws_auth import AWSRequestsAuth
import requests
import _snowflake
import json

def main_handler(name):
    # シークレットから一時的な認証情報を取得
    cloud_provider_object = _snowflake.get_cloud_provider_token('cred')

    # AWSRequestsAuthに直接資格情報を渡す
    auth = AWSRequestsAuth(
        aws_access_key=cloud_provider_object.access_key_id,
        aws_secret_access_key=cloud_provider_object.secret_access_key,
        aws_token=cloud_provider_object.token,
        aws_host='{api_id}.execute-api.ap-northeast-1.amazonaws.com',
        aws_region='ap-northeast-1',
        aws_service='execute-api'
    )

    # nameパラメータでペイロードを定義
    payload = {'name': name}

    # API GatewayのエンドポイントにPOSTリクエストを送信
    url = 'https://{api_id}.execute-api.ap-northeast-1.amazonaws.com/dev/hello'
    response = requests.post(url, json=payload, auth=auth)

    # JSONレスポンスをパースしてbodyを抽出
    response_data = json.loads(response.text)
    body_content = json.loads(response_data["body"])

    return body_content
$$;

>SELECT aws_api_gateway_post_function('Alice');
+----------------------------------------+
| AWS_API_GATEWAY_POST_FUNCTION('ALICE') |
|----------------------------------------|
| Hello, Alice                           |
+----------------------------------------+
1 Row(s) produced. Time Elapsed: 3.982s

>SELECT aws_api_gateway_post_function('John');
+---------------------------------------+
| AWS_API_GATEWAY_POST_FUNCTION('JOHN') |
|---------------------------------------|
| Hello, John                           |
+---------------------------------------+
1 Row(s) produced. Time Elapsed: 2.101s